{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neo4j Graph Store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"API_KEY_HERE\"\n",
    "\n",
    "import logging\n",
    "import sys\n",
    "from llama_index.llms import OpenAI\n",
    "from llama_index import ServiceContext\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "\n",
    "# define LLM\n",
    "llm = OpenAI(temperature=0, model=\"gpt-3.5-turbo\")\n",
    "service_context = ServiceContext.from_defaults(llm=llm, chunk_size=512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "'gpt-3.5-turbo' is not a valid OpenAIEmbeddingModelType",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[17], line 43\u001b[0m\n\u001b[1;32m     30\u001b[0m llm \u001b[38;5;241m=\u001b[39m AzureOpenAI(\n\u001b[1;32m     31\u001b[0m     deployment_name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mgpt-35-turbo-16k\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     32\u001b[0m     temperature\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     39\u001b[0m     },\n\u001b[1;32m     40\u001b[0m )\n\u001b[1;32m     42\u001b[0m \u001b[38;5;66;03m# You need to deploy your own embedding model as well as your own chat completion model\u001b[39;00m\n\u001b[0;32m---> 43\u001b[0m embedding_llm \u001b[38;5;241m=\u001b[39m \u001b[43mOpenAIEmbedding\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     44\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mgpt-3.5-turbo\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     45\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdeployment_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mgpt-35-turbo-16k\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     46\u001b[0m \u001b[43m    \u001b[49m\u001b[43mapi_key\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mopenai\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapi_key\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     47\u001b[0m \u001b[43m    \u001b[49m\u001b[43mapi_base\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mopenai\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapi_base\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     48\u001b[0m \u001b[43m    \u001b[49m\u001b[43mapi_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mopenai\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapi_type\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     49\u001b[0m \u001b[43m    \u001b[49m\u001b[43mapi_version\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mopenai\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapi_version\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     50\u001b[0m \u001b[43m)\u001b[49m\n\u001b[1;32m     52\u001b[0m service_context \u001b[38;5;241m=\u001b[39m ServiceContext\u001b[38;5;241m.\u001b[39mfrom_defaults(\n\u001b[1;32m     53\u001b[0m     llm\u001b[38;5;241m=\u001b[39mllm,\n\u001b[1;32m     54\u001b[0m     embed_model\u001b[38;5;241m=\u001b[39membedding_llm,\n\u001b[1;32m     55\u001b[0m )\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/llama_index/embeddings/openai.py:280\u001b[0m, in \u001b[0;36mOpenAIEmbedding.__init__\u001b[0;34m(self, mode, model, embed_batch_size, additional_kwargs, api_key, api_base, api_version, max_retries, timeout, reuse_client, callback_manager, default_headers, http_client, **kwargs)\u001b[0m\n\u001b[1;32m    272\u001b[0m additional_kwargs \u001b[38;5;241m=\u001b[39m additional_kwargs \u001b[38;5;129;01mor\u001b[39;00m {}\n\u001b[1;32m    274\u001b[0m api_key, api_base, api_version \u001b[38;5;241m=\u001b[39m resolve_openai_credentials(\n\u001b[1;32m    275\u001b[0m     api_key\u001b[38;5;241m=\u001b[39mapi_key,\n\u001b[1;32m    276\u001b[0m     api_base\u001b[38;5;241m=\u001b[39mapi_base,\n\u001b[1;32m    277\u001b[0m     api_version\u001b[38;5;241m=\u001b[39mapi_version,\n\u001b[1;32m    278\u001b[0m )\n\u001b[0;32m--> 280\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_query_engine \u001b[38;5;241m=\u001b[39m \u001b[43mget_engine\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m_QUERY_MODE_MODEL_DICT\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    281\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_text_engine \u001b[38;5;241m=\u001b[39m get_engine(mode, model, _TEXT_MODE_MODEL_DICT)\n\u001b[1;32m    283\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodel_name\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m kwargs:\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/llama_index/embeddings/openai.py:198\u001b[0m, in \u001b[0;36mget_engine\u001b[0;34m(mode, model, mode_model_dict)\u001b[0m\n\u001b[1;32m    192\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mget_engine\u001b[39m(\n\u001b[1;32m    193\u001b[0m     mode: \u001b[38;5;28mstr\u001b[39m,\n\u001b[1;32m    194\u001b[0m     model: \u001b[38;5;28mstr\u001b[39m,\n\u001b[1;32m    195\u001b[0m     mode_model_dict: Dict[Tuple[OpenAIEmbeddingMode, \u001b[38;5;28mstr\u001b[39m], OpenAIEmbeddingModeModel],\n\u001b[1;32m    196\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m OpenAIEmbeddingModeModel:\n\u001b[1;32m    197\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Get engine.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 198\u001b[0m     key \u001b[38;5;241m=\u001b[39m (OpenAIEmbeddingMode(mode), \u001b[43mOpenAIEmbeddingModelType\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m    199\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m key \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m mode_model_dict:\n\u001b[1;32m    200\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mInvalid mode, model combination: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mkey\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/enum.py:740\u001b[0m, in \u001b[0;36mEnumType.__call__\u001b[0;34m(cls, value, names, module, qualname, type, start, boundary, *values)\u001b[0m\n\u001b[1;32m    738\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m names:\n\u001b[1;32m    739\u001b[0m         value \u001b[38;5;241m=\u001b[39m (value, names) \u001b[38;5;241m+\u001b[39m values\n\u001b[0;32m--> 740\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__new__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mcls\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    741\u001b[0m \u001b[38;5;66;03m# otherwise, functional API: we're creating a new Enum type\u001b[39;00m\n\u001b[1;32m    742\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m names \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mtype\u001b[39m \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    743\u001b[0m     \u001b[38;5;66;03m# no body? no data-type? possibly wrong usage\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/enum.py:1154\u001b[0m, in \u001b[0;36mEnum.__new__\u001b[0;34m(cls, value)\u001b[0m\n\u001b[1;32m   1152\u001b[0m ve_exc \u001b[38;5;241m=\u001b[39m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m%r\u001b[39;00m\u001b[38;5;124m is not a valid \u001b[39m\u001b[38;5;132;01m%s\u001b[39;00m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m%\u001b[39m (value, \u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__qualname__\u001b[39m))\n\u001b[1;32m   1153\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m result \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m exc \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m-> 1154\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m ve_exc\n\u001b[1;32m   1155\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m exc \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m   1156\u001b[0m     exc \u001b[38;5;241m=\u001b[39m \u001b[38;5;167;01mTypeError\u001b[39;00m(\n\u001b[1;32m   1157\u001b[0m             \u001b[38;5;124m'\u001b[39m\u001b[38;5;124merror in \u001b[39m\u001b[38;5;132;01m%s\u001b[39;00m\u001b[38;5;124m._missing_: returned \u001b[39m\u001b[38;5;132;01m%r\u001b[39;00m\u001b[38;5;124m instead of None or a valid member\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m   1158\u001b[0m             \u001b[38;5;241m%\u001b[39m (\u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, result)\n\u001b[1;32m   1159\u001b[0m             )\n",
      "\u001b[0;31mValueError\u001b[0m: 'gpt-3.5-turbo' is not a valid OpenAIEmbeddingModelType"
     ]
    }
   ],
   "source": [
    "# For Azure OpenAI\n",
    "import os\n",
    "import json\n",
    "import openai\n",
    "from llama_index.llms import AzureOpenAI\n",
    "from llama_index.embeddings import OpenAIEmbedding\n",
    "from llama_index import (\n",
    "    VectorStoreIndex,\n",
    "    SimpleDirectoryReader,\n",
    "    KnowledgeGraphIndex,\n",
    "    ServiceContext,\n",
    ")\n",
    "\n",
    "import logging\n",
    "import sys\n",
    "\n",
    "from IPython.display import Markdown, display\n",
    "\n",
    "logging.basicConfig(\n",
    "    stream=sys.stdout, level=logging.INFO\n",
    ")  # logging.DEBUG for more verbose output\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "\n",
    "openai.api_type = \"azure\"\n",
    "openai.api_base = \"https://firstaiproject.openai.azure.com\"\n",
    "openai.api_version = \"2022-12-01\"\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"77a8bace65c148c7a10be35741adaf5b\"\n",
    "openai.api_key = os.getenv(\"OPENAI_API_KEY\")\n",
    "\n",
    "llm = AzureOpenAI(\n",
    "    deployment_name=\"gpt-35-turbo-16k\",\n",
    "    temperature=0,\n",
    "    openai_api_version=openai.api_version,\n",
    "    model_kwargs={\n",
    "        \"api_key\": openai.api_key,\n",
    "        \"api_base\": openai.api_base,\n",
    "        \"api_type\": openai.api_type,\n",
    "        \"api_version\": openai.api_version,\n",
    "    },\n",
    ")\n",
    "\n",
    "# You need to deploy your own embedding model as well as your own chat completion model\n",
    "embedding_llm = OpenAIEmbedding(\n",
    "    model=\"text-embedding-ada-002\",\n",
    "    deployment_name=\"gpt-35-turbo-16k\",\n",
    "    api_key=openai.api_key,\n",
    "    api_base=openai.api_base,\n",
    "    api_type=openai.api_type,\n",
    "    api_version=openai.api_version,\n",
    ")\n",
    "\n",
    "service_context = ServiceContext.from_defaults(\n",
    "    llm=llm,\n",
    "    embed_model=embedding_llm,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Knowledge Graph with Neo4jGraphStore\n",
    "### Building the Knowledge Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import (\n",
    "    KnowledgeGraphIndex,\n",
    "    ServiceContext,\n",
    "    SimpleDirectoryReader,\n",
    ")\n",
    "from llama_index.storage.storage_context import StorageContext\n",
    "from llama_index.graph_stores import Neo4jGraphStore\n",
    "\n",
    "\n",
    "from llama_index.llms import OpenAI\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = SimpleDirectoryReader(\"./data/paul_graham\").load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare for Neo4j"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 安装 neo4j 服务\n",
    "```shell\n",
    "docker run -d -p 7474:7474 -p 7687:7687 --name neo4j-apoc -e NEO4J_apoc_export_file_enabled=true -e NEO4J_apoc_import_file_enabled=true -e NEO4J_apoc_import_file_use__neo4j__config=true -e NEO4J_AUTH=neo4j/pleaseletmein -e NEO4J_PLUGINS=\\[\\\"apoc\\\"\\] neo4j:latest\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: neo4j in /Users/minp/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages (5.15.0)\n",
      "Requirement already satisfied: pytz in /Users/minp/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages (from neo4j) (2023.3.post1)\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install neo4j\n",
    "\n",
    "username = \"neo4j\"\n",
    "password = \"pleaseletmein\"\n",
    "url = \"bolt://localhost:7687\"\n",
    "database = \"neo4j\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate Neo4jGraph KG Indexes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://firstaiproject.openai.azure.com//openai/deployments/text-embedding-ada-002/chat/completions?api-version=2023-06-01-preview \"HTTP/1.1 400 Bad Request\"\n",
      "HTTP Request: POST https://firstaiproject.openai.azure.com//openai/deployments/text-embedding-ada-002/chat/completions?api-version=2023-06-01-preview \"HTTP/1.1 400 Bad Request\"\n",
      "HTTP Request: POST https://firstaiproject.openai.azure.com//openai/deployments/text-embedding-ada-002/chat/completions?api-version=2023-06-01-preview \"HTTP/1.1 400 Bad Request\"\n",
      "HTTP Request: POST https://firstaiproject.openai.azure.com//openai/deployments/text-embedding-ada-002/chat/completions?api-version=2023-06-01-preview \"HTTP/1.1 400 Bad Request\"\n"
     ]
    },
    {
     "ename": "BadRequestError",
     "evalue": "Error code: 400 - {'error': {'code': 'OperationNotSupported', 'message': 'The chatCompletion operation does not work with the specified model, text-embedding-ada-002. Please choose different model and try again. You can learn more about which models can be used with each operation here: https://go.microsoft.com/fwlink/?linkid=2197993.'}}",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mBadRequestError\u001b[0m                           Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[14], line 11\u001b[0m\n\u001b[1;32m      8\u001b[0m storage_context \u001b[38;5;241m=\u001b[39m StorageContext\u001b[38;5;241m.\u001b[39mfrom_defaults(graph_store\u001b[38;5;241m=\u001b[39mgraph_store)\n\u001b[1;32m     10\u001b[0m \u001b[38;5;66;03m# NOTE: can take a while!\u001b[39;00m\n\u001b[0;32m---> 11\u001b[0m index \u001b[38;5;241m=\u001b[39m \u001b[43mKnowledgeGraphIndex\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_documents\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     12\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdocuments\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     13\u001b[0m \u001b[43m    \u001b[49m\u001b[43mstorage_context\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstorage_context\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     14\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmax_triplets_per_chunk\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     15\u001b[0m \u001b[43m    \u001b[49m\u001b[43mservice_context\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mservice_context\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     16\u001b[0m \u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/llama_index/indices/base.py:106\u001b[0m, in \u001b[0;36mBaseIndex.from_documents\u001b[0;34m(cls, documents, storage_context, service_context, show_progress, **kwargs)\u001b[0m\n\u001b[1;32m     97\u001b[0m     docstore\u001b[38;5;241m.\u001b[39mset_document_hash(doc\u001b[38;5;241m.\u001b[39mget_doc_id(), doc\u001b[38;5;241m.\u001b[39mhash)\n\u001b[1;32m     99\u001b[0m nodes \u001b[38;5;241m=\u001b[39m run_transformations(\n\u001b[1;32m    100\u001b[0m     documents,  \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[1;32m    101\u001b[0m     service_context\u001b[38;5;241m.\u001b[39mtransformations,\n\u001b[1;32m    102\u001b[0m     show_progress\u001b[38;5;241m=\u001b[39mshow_progress,\n\u001b[1;32m    103\u001b[0m     \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs,\n\u001b[1;32m    104\u001b[0m )\n\u001b[0;32m--> 106\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m    107\u001b[0m \u001b[43m    \u001b[49m\u001b[43mnodes\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mnodes\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    108\u001b[0m \u001b[43m    \u001b[49m\u001b[43mstorage_context\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstorage_context\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    109\u001b[0m \u001b[43m    \u001b[49m\u001b[43mservice_context\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mservice_context\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    110\u001b[0m \u001b[43m    \u001b[49m\u001b[43mshow_progress\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mshow_progress\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    111\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    112\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/llama_index/indices/knowledge_graph/base.py:81\u001b[0m, in \u001b[0;36mKnowledgeGraphIndex.__init__\u001b[0;34m(self, nodes, index_struct, service_context, storage_context, kg_triple_extract_template, max_triplets_per_chunk, include_embeddings, show_progress, max_object_length, kg_triplet_extract_fn, **kwargs)\u001b[0m\n\u001b[1;32m     78\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_max_object_length \u001b[38;5;241m=\u001b[39m max_object_length\n\u001b[1;32m     79\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_kg_triplet_extract_fn \u001b[38;5;241m=\u001b[39m kg_triplet_extract_fn\n\u001b[0;32m---> 81\u001b[0m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m     82\u001b[0m \u001b[43m    \u001b[49m\u001b[43mnodes\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mnodes\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     83\u001b[0m \u001b[43m    \u001b[49m\u001b[43mindex_struct\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mindex_struct\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     84\u001b[0m \u001b[43m    \u001b[49m\u001b[43mservice_context\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mservice_context\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     85\u001b[0m \u001b[43m    \u001b[49m\u001b[43mstorage_context\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstorage_context\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     86\u001b[0m \u001b[43m    \u001b[49m\u001b[43mshow_progress\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mshow_progress\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     87\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     88\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     90\u001b[0m \u001b[38;5;66;03m# TODO: legacy conversion - remove in next release\u001b[39;00m\n\u001b[1;32m     91\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[1;32m     92\u001b[0m     \u001b[38;5;28mlen\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mindex_struct\u001b[38;5;241m.\u001b[39mtable) \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m     93\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgraph_store, SimpleGraphStore)\n\u001b[1;32m     94\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgraph_store\u001b[38;5;241m.\u001b[39m_data\u001b[38;5;241m.\u001b[39mgraph_dict) \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m     95\u001b[0m ):\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/llama_index/indices/base.py:71\u001b[0m, in \u001b[0;36mBaseIndex.__init__\u001b[0;34m(self, nodes, index_struct, storage_context, service_context, show_progress, **kwargs)\u001b[0m\n\u001b[1;32m     69\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m index_struct \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m     70\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m nodes \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m---> 71\u001b[0m     index_struct \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbuild_index_from_nodes\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnodes\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     72\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_index_struct \u001b[38;5;241m=\u001b[39m index_struct\n\u001b[1;32m     73\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_storage_context\u001b[38;5;241m.\u001b[39mindex_store\u001b[38;5;241m.\u001b[39madd_index_struct(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_index_struct)\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/llama_index/indices/base.py:175\u001b[0m, in \u001b[0;36mBaseIndex.build_index_from_nodes\u001b[0;34m(self, nodes)\u001b[0m\n\u001b[1;32m    173\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Build the index from nodes.\"\"\"\u001b[39;00m\n\u001b[1;32m    174\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_docstore\u001b[38;5;241m.\u001b[39madd_documents(nodes, allow_update\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m--> 175\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_build_index_from_nodes\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnodes\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/llama_index/indices/knowledge_graph/base.py:167\u001b[0m, in \u001b[0;36mKnowledgeGraphIndex._build_index_from_nodes\u001b[0;34m(self, nodes)\u001b[0m\n\u001b[1;32m    163\u001b[0m nodes_with_progress \u001b[38;5;241m=\u001b[39m get_tqdm_iterable(\n\u001b[1;32m    164\u001b[0m     nodes, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_show_progress, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mProcessing nodes\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    165\u001b[0m )\n\u001b[1;32m    166\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m n \u001b[38;5;129;01min\u001b[39;00m nodes_with_progress:\n\u001b[0;32m--> 167\u001b[0m     triplets \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_extract_triplets\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    168\u001b[0m \u001b[43m        \u001b[49m\u001b[43mn\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_content\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmetadata_mode\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mMetadataMode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mLLM\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    169\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    170\u001b[0m     logger\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m> Extracted triplets: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtriplets\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    171\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m triplet \u001b[38;5;129;01min\u001b[39;00m triplets:\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/llama_index/indices/knowledge_graph/base.py:118\u001b[0m, in \u001b[0;36mKnowledgeGraphIndex._extract_triplets\u001b[0;34m(self, text)\u001b[0m\n\u001b[1;32m    116\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_kg_triplet_extract_fn(text)\n\u001b[1;32m    117\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 118\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_llm_extract_triplets\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtext\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/llama_index/indices/knowledge_graph/base.py:122\u001b[0m, in \u001b[0;36mKnowledgeGraphIndex._llm_extract_triplets\u001b[0;34m(self, text)\u001b[0m\n\u001b[1;32m    120\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_llm_extract_triplets\u001b[39m(\u001b[38;5;28mself\u001b[39m, text: \u001b[38;5;28mstr\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m List[Tuple[\u001b[38;5;28mstr\u001b[39m, \u001b[38;5;28mstr\u001b[39m, \u001b[38;5;28mstr\u001b[39m]]:\n\u001b[1;32m    121\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Extract keywords from text.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 122\u001b[0m     response \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_service_context\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mllm\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpredict\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    123\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkg_triple_extract_template\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    124\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtext\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtext\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    125\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    126\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_parse_triplet_response(\n\u001b[1;32m    127\u001b[0m         response, max_length\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_max_object_length\n\u001b[1;32m    128\u001b[0m     )\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/llama_index/llms/llm.py:220\u001b[0m, in \u001b[0;36mLLM.predict\u001b[0;34m(self, prompt, **prompt_args)\u001b[0m\n\u001b[1;32m    218\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmetadata\u001b[38;5;241m.\u001b[39mis_chat_model:\n\u001b[1;32m    219\u001b[0m     messages \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_get_messages(prompt, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mprompt_args)\n\u001b[0;32m--> 220\u001b[0m     chat_response \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mchat\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmessages\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    221\u001b[0m     output \u001b[38;5;241m=\u001b[39m chat_response\u001b[38;5;241m.\u001b[39mmessage\u001b[38;5;241m.\u001b[39mcontent \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    222\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/llama_index/llms/base.py:97\u001b[0m, in \u001b[0;36mllm_chat_callback.<locals>.wrap.<locals>.wrapped_llm_chat\u001b[0;34m(_self, messages, **kwargs)\u001b[0m\n\u001b[1;32m     88\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m wrapper_logic(_self) \u001b[38;5;28;01mas\u001b[39;00m callback_manager:\n\u001b[1;32m     89\u001b[0m     event_id \u001b[38;5;241m=\u001b[39m callback_manager\u001b[38;5;241m.\u001b[39mon_event_start(\n\u001b[1;32m     90\u001b[0m         CBEventType\u001b[38;5;241m.\u001b[39mLLM,\n\u001b[1;32m     91\u001b[0m         payload\u001b[38;5;241m=\u001b[39m{\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     95\u001b[0m         },\n\u001b[1;32m     96\u001b[0m     )\n\u001b[0;32m---> 97\u001b[0m     f_return_val \u001b[38;5;241m=\u001b[39m \u001b[43mf\u001b[49m\u001b[43m(\u001b[49m\u001b[43m_self\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmessages\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     99\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(f_return_val, Generator):\n\u001b[1;32m    100\u001b[0m         \u001b[38;5;66;03m# intercept the generator and add a callback to the end\u001b[39;00m\n\u001b[1;32m    101\u001b[0m         \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mwrapped_gen\u001b[39m() \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m ChatResponseGen:\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/llama_index/llms/openai.py:234\u001b[0m, in \u001b[0;36mOpenAI.chat\u001b[0;34m(self, messages, **kwargs)\u001b[0m\n\u001b[1;32m    232\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    233\u001b[0m     chat_fn \u001b[38;5;241m=\u001b[39m completion_to_chat_decorator(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_complete)\n\u001b[0;32m--> 234\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mchat_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmessages\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/llama_index/llms/openai.py:289\u001b[0m, in \u001b[0;36mOpenAI._chat\u001b[0;34m(self, messages, **kwargs)\u001b[0m\n\u001b[1;32m    287\u001b[0m client \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_get_client()\n\u001b[1;32m    288\u001b[0m message_dicts \u001b[38;5;241m=\u001b[39m to_openai_message_dicts(messages)\n\u001b[0;32m--> 289\u001b[0m response \u001b[38;5;241m=\u001b[39m \u001b[43mclient\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mchat\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcompletions\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcreate\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    290\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmessages\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmessage_dicts\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    291\u001b[0m \u001b[43m    \u001b[49m\u001b[43mstream\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    292\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_get_model_kwargs\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    293\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    294\u001b[0m openai_message \u001b[38;5;241m=\u001b[39m response\u001b[38;5;241m.\u001b[39mchoices[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mmessage\n\u001b[1;32m    295\u001b[0m message \u001b[38;5;241m=\u001b[39m from_openai_message(openai_message)\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/openai/_utils/_utils.py:270\u001b[0m, in \u001b[0;36mrequired_args.<locals>.inner.<locals>.wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    268\u001b[0m             msg \u001b[38;5;241m=\u001b[39m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMissing required argument: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mquote(missing[\u001b[38;5;241m0\u001b[39m])\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    269\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(msg)\n\u001b[0;32m--> 270\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/openai/resources/chat/completions.py:645\u001b[0m, in \u001b[0;36mCompletions.create\u001b[0;34m(self, messages, model, frequency_penalty, function_call, functions, logit_bias, logprobs, max_tokens, n, presence_penalty, response_format, seed, stop, stream, temperature, tool_choice, tools, top_logprobs, top_p, user, extra_headers, extra_query, extra_body, timeout)\u001b[0m\n\u001b[1;32m    596\u001b[0m \u001b[38;5;129m@required_args\u001b[39m([\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmessages\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodel\u001b[39m\u001b[38;5;124m\"\u001b[39m], [\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmessages\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodel\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstream\u001b[39m\u001b[38;5;124m\"\u001b[39m])\n\u001b[1;32m    597\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcreate\u001b[39m(\n\u001b[1;32m    598\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    643\u001b[0m     timeout: \u001b[38;5;28mfloat\u001b[39m \u001b[38;5;241m|\u001b[39m httpx\u001b[38;5;241m.\u001b[39mTimeout \u001b[38;5;241m|\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;241m|\u001b[39m NotGiven \u001b[38;5;241m=\u001b[39m NOT_GIVEN,\n\u001b[1;32m    644\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m ChatCompletion \u001b[38;5;241m|\u001b[39m Stream[ChatCompletionChunk]:\n\u001b[0;32m--> 645\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_post\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    646\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m/chat/completions\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    647\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbody\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmaybe_transform\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    648\u001b[0m \u001b[43m            \u001b[49m\u001b[43m{\u001b[49m\n\u001b[1;32m    649\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mmessages\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mmessages\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    650\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mmodel\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    651\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mfrequency_penalty\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mfrequency_penalty\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    652\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mfunction_call\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mfunction_call\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    653\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mfunctions\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mfunctions\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    654\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mlogit_bias\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mlogit_bias\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    655\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mlogprobs\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mlogprobs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    656\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mmax_tokens\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mmax_tokens\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    657\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mn\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mn\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    658\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mpresence_penalty\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mpresence_penalty\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    659\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mresponse_format\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mresponse_format\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    660\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mseed\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mseed\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    661\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mstop\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mstop\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    662\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mstream\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mstream\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    663\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mtemperature\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mtemperature\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    664\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mtool_choice\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mtool_choice\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    665\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mtools\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mtools\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    666\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mtop_logprobs\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mtop_logprobs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    667\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mtop_p\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mtop_p\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    668\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43muser\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43muser\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    669\u001b[0m \u001b[43m            \u001b[49m\u001b[43m}\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    670\u001b[0m \u001b[43m            \u001b[49m\u001b[43mcompletion_create_params\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mCompletionCreateParams\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    671\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    672\u001b[0m \u001b[43m        \u001b[49m\u001b[43moptions\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmake_request_options\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    673\u001b[0m \u001b[43m            \u001b[49m\u001b[43mextra_headers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mextra_headers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mextra_query\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mextra_query\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mextra_body\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mextra_body\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\n\u001b[1;32m    674\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    675\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcast_to\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mChatCompletion\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    676\u001b[0m \u001b[43m        \u001b[49m\u001b[43mstream\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstream\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    677\u001b[0m \u001b[43m        \u001b[49m\u001b[43mstream_cls\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mStream\u001b[49m\u001b[43m[\u001b[49m\u001b[43mChatCompletionChunk\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    678\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/openai/_base_client.py:1088\u001b[0m, in \u001b[0;36mSyncAPIClient.post\u001b[0;34m(self, path, cast_to, body, options, files, stream, stream_cls)\u001b[0m\n\u001b[1;32m   1074\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mpost\u001b[39m(\n\u001b[1;32m   1075\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m   1076\u001b[0m     path: \u001b[38;5;28mstr\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1083\u001b[0m     stream_cls: \u001b[38;5;28mtype\u001b[39m[_StreamT] \u001b[38;5;241m|\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m   1084\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m ResponseT \u001b[38;5;241m|\u001b[39m _StreamT:\n\u001b[1;32m   1085\u001b[0m     opts \u001b[38;5;241m=\u001b[39m FinalRequestOptions\u001b[38;5;241m.\u001b[39mconstruct(\n\u001b[1;32m   1086\u001b[0m         method\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpost\u001b[39m\u001b[38;5;124m\"\u001b[39m, url\u001b[38;5;241m=\u001b[39mpath, json_data\u001b[38;5;241m=\u001b[39mbody, files\u001b[38;5;241m=\u001b[39mto_httpx_files(files), \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39moptions\n\u001b[1;32m   1087\u001b[0m     )\n\u001b[0;32m-> 1088\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m cast(ResponseT, \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrequest\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcast_to\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mopts\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstream\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstream\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstream_cls\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstream_cls\u001b[49m\u001b[43m)\u001b[49m)\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/openai/_base_client.py:853\u001b[0m, in \u001b[0;36mSyncAPIClient.request\u001b[0;34m(self, cast_to, options, remaining_retries, stream, stream_cls)\u001b[0m\n\u001b[1;32m    844\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mrequest\u001b[39m(\n\u001b[1;32m    845\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m    846\u001b[0m     cast_to: Type[ResponseT],\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    851\u001b[0m     stream_cls: \u001b[38;5;28mtype\u001b[39m[_StreamT] \u001b[38;5;241m|\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    852\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m ResponseT \u001b[38;5;241m|\u001b[39m _StreamT:\n\u001b[0;32m--> 853\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_request\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    854\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcast_to\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcast_to\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    855\u001b[0m \u001b[43m        \u001b[49m\u001b[43moptions\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moptions\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    856\u001b[0m \u001b[43m        \u001b[49m\u001b[43mstream\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstream\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    857\u001b[0m \u001b[43m        \u001b[49m\u001b[43mstream_cls\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstream_cls\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    858\u001b[0m \u001b[43m        \u001b[49m\u001b[43mremaining_retries\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mremaining_retries\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    859\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/llm_notebooks/lib/python3.12/site-packages/openai/_base_client.py:930\u001b[0m, in \u001b[0;36mSyncAPIClient._request\u001b[0;34m(self, cast_to, options, remaining_retries, stream, stream_cls)\u001b[0m\n\u001b[1;32m    927\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m err\u001b[38;5;241m.\u001b[39mresponse\u001b[38;5;241m.\u001b[39mis_closed:\n\u001b[1;32m    928\u001b[0m         err\u001b[38;5;241m.\u001b[39mresponse\u001b[38;5;241m.\u001b[39mread()\n\u001b[0;32m--> 930\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_make_status_error_from_response(err\u001b[38;5;241m.\u001b[39mresponse) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    932\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_process_response(\n\u001b[1;32m    933\u001b[0m     cast_to\u001b[38;5;241m=\u001b[39mcast_to,\n\u001b[1;32m    934\u001b[0m     options\u001b[38;5;241m=\u001b[39moptions,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    937\u001b[0m     stream_cls\u001b[38;5;241m=\u001b[39mstream_cls,\n\u001b[1;32m    938\u001b[0m )\n",
      "\u001b[0;31mBadRequestError\u001b[0m: Error code: 400 - {'error': {'code': 'OperationNotSupported', 'message': 'The chatCompletion operation does not work with the specified model, text-embedding-ada-002. Please choose different model and try again. You can learn more about which models can be used with each operation here: https://go.microsoft.com/fwlink/?linkid=2197993.'}}"
     ]
    }
   ],
   "source": [
    "graph_store = Neo4jGraphStore(\n",
    "    username=username,\n",
    "    password=password,\n",
    "    url=url,\n",
    "    database=database,\n",
    ")\n",
    "\n",
    "storage_context = StorageContext.from_defaults(graph_store=graph_store)\n",
    "\n",
    "# NOTE: can take a while!\n",
    "index = KnowledgeGraphIndex.from_documents(\n",
    "    documents,\n",
    "    storage_context=storage_context,\n",
    "    max_triplets_per_chunk=2,\n",
    "    service_context=service_context,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Querying the Knowledge Graph\n",
    "First, we can query and send only the triplets to the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    include_text=False, response_mode=\"tree_summarize\"\n",
    ")\n",
    "\n",
    "response = query_engine.query(\"Tell me more about Interleaf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(Markdown(f\"<b>{response}</b>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    include_text=True, response_mode=\"tree_summarize\"\n",
    ")\n",
    "response = query_engine.query(\n",
    "    \"Tell me more about what the author worked on at Interleaf\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(Markdown(f\"<b>{response}</b>\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Query with embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean dataset first\n",
    "graph_store.query(\n",
    "    \"\"\"\n",
    "MATCH (n) DETACH DELETE n\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "# NOTE: can take a while!\n",
    "index = KnowledgeGraphIndex.from_documents(\n",
    "    documents,\n",
    "    storage_context=storage_context,\n",
    "    max_triplets_per_chunk=2,\n",
    "    service_context=service_context,\n",
    "    include_embeddings=True,\n",
    ")\n",
    "\n",
    "query_engine = index.as_query_engine(\n",
    "    include_text=True,\n",
    "    response_mode=\"tree_summarize\",\n",
    "    embedding_mode=\"hybrid\",\n",
    "    similarity_top_k=5,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# query using top 3 triplets plus keywords (duplicate triplets are removed)\n",
    "response = query_engine.query(\n",
    "    \"Tell me more about what the author worked on at Interleaf\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(Markdown(f\"<b>{response}</b>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.node_parser import SentenceSplitter\n",
    "node_parser = SentenceSplitter()\n",
    "nodes = node_parser.get_nodes_from_documents(documents)\n",
    "# initialize an empty index for now\n",
    "index = KnowledgeGraphIndex.from_documents([], storage_context=storage_context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add keyword mappings and nodes manually\n",
    "# add triplets (subject, relationship, object)\n",
    "\n",
    "# for node 0\n",
    "node_0_tups = [\n",
    "    (\"author\", \"worked on\", \"writing\"),\n",
    "    (\"author\", \"worked on\", \"programming\"),\n",
    "]\n",
    "for tup in node_0_tups:\n",
    "    index.upsert_triplet_and_node(tup, nodes[0])\n",
    "\n",
    "# for node 1\n",
    "node_1_tups = [\n",
    "    (\"Interleaf\", \"made software for\", \"creating documents\"),\n",
    "    (\"Interleaf\", \"added\", \"scripting language\"),\n",
    "    (\"software\", \"generate\", \"web sites\"),\n",
    "]\n",
    "for tup in node_1_tups:\n",
    "    index.upsert_triplet_and_node(tup, nodes[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    include_text=False, response_mode=\"tree_summarize\"\n",
    ")\n",
    "\n",
    "response = query_engine.query(\"Tell me more about Interleaf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(Markdown(f\"<b>{response}</b>\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
